//
// Created by Never on 2019/11/1.
//

#ifndef MINIMALSPANNINGFOREST_MINIMALSPANNINGFORESTWITHCAPCONS_H
#define MINIMALSPANNINGFOREST_MINIMALSPANNINGFORESTWITHCAPCONS_H

#include <cmath>

#include "MinimalSpanningForest.h"
template <typename val>
class MinimalSpanningForestWithCapCons: public MinimalSpanningForest<val>{
private:
    std::vector<int> demands;
    int vehicleCap;
    std::vector<int> demandsInBranch;
    std::vector<int> demandsInRoots;
    std::vector<int> roots;
    std::vector<int> preRoot;
    std::vector<int> firstNodeInBranches;
    std::vector<int> partialRoots;
    const int unable_to_insert = INT32_MAX;
    bool isAbort;

    void initialTableToDefault() override ;
    void initialTable(std::vector<int> usedRoots) override ;
    bool isKnownNodeToNodeCloser(int knownNode, int node) override ;
    bool isKnownNodeLinkNodeSatisfyCapCons(int knownNode, int node);
    bool isSolved();
    void handleNewInsertNode(int node) override ;
    void updateMinCostAfterInsertingNode(int node) override ;
    bool findClosestKnownNodeToNode(int node);
    bool isExit(int node) override ;
public:
    MinimalSpanningForestWithCapCons(const std::vector<int>& roots, const std::vector<int>& nodes,
                                     std::vector<std::vector<val>> distances,
                                     std::vector<int> demands, int vehicleCap)
            :
            MinimalSpanningForest<val>(roots, nodes, distances), demands(std::move(demands)), vehicleCap(vehicleCap),
            demandsInBranch(nodes.size()+roots.size(), 0),
            demandsInRoots(roots.size()+nodes.size(), 0), roots(roots),
            preRoot(roots.size()+nodes.size()), firstNodeInBranches(roots.size() + nodes.size())
            {};
    bool minimalSpanningForestWithCapCons(const std::vector<int>& usedRoots);
    std::map<int, int> getDepotsToVehicleNumber();
    std::map<int, std::vector<int>> getNodesInEachRoot();
};

template<typename val>
void MinimalSpanningForestWithCapCons<val>::initialTableToDefault() {
    MinimalSpanningForest<val>::initialTableToDefault();
    for(int i = 0; i < demandsInRoots.size(); i++)
    {
        demandsInRoots[i] = 0;
        demandsInBranch[i] = 0;
    }
}

template<typename val>
void MinimalSpanningForestWithCapCons<val>::initialTable(std::vector<int> usedRoots) {
    MinimalSpanningForest<val>::initialTable(usedRoots);
    for(auto node: MinimalSpanningForest<val>::nodes)
        preRoot[node] = firstNodeInBranches[node] = MinimalSpanningForest<val>::preNode[node];
    for(auto root: roots)
        preRoot[root] = firstNodeInBranches[root] = MinimalSpanningForest<val>::preNode[root];
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::minimalSpanningForestWithCapCons(const std::vector<int>& usedRoots) {
    partialRoots = usedRoots;
    isAbort = false;
    MinimalSpanningForest<val>::minimalSpanningForest(usedRoots);
    return isSolved();
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::isKnownNodeToNodeCloser(int knownNode, int node) {
    return isKnownNodeLinkNodeSatisfyCapCons(knownNode, node) and
    MinimalSpanningForest<val>::isKnownNodeToNodeCloser(knownNode, node);
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::isKnownNodeLinkNodeSatisfyCapCons(int knownNode, int node) {
    int root = preRoot[knownNode];
    return demandsInRoots[root] + demands[node] <= demands[root];
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::isSolved() {
    for(auto node: MinimalSpanningForest<val>::nodes)
    {
        if(not MinimalSpanningForest<val>::isKnown[node])
            return false;
    }
    return true;
}

template<typename val>
void MinimalSpanningForestWithCapCons<val>::handleNewInsertNode(int node) {
    int preNode = MinimalSpanningForest<val>::preNode[node];
    int root = preRoot[preNode], firstNodeInBranch = firstNodeInBranches[preNode];
    if(firstNodeInBranch == preNode)
        firstNodeInBranch = node;
    preRoot[node] = root;
    firstNodeInBranches[node] = firstNodeInBranch;
    demandsInRoots[root] += demands[node];
    demandsInBranch[firstNodeInBranch] += demands[node];
    MinimalSpanningForest<val>::handleNewInsertNode(node);
}

template<typename val>
void MinimalSpanningForestWithCapCons<val>::updateMinCostAfterInsertingNode(int node) {
    for(auto n: MinimalSpanningForest<val>::nodes)
    {
        if(not MinimalSpanningForest<val>::isKnown[n])
        {
            int root = preRoot[MinimalSpanningForest<val>::preNode[n]];
            if(demandsInRoots[root]+demands[n] > demands[root])
            {
                MinimalSpanningForest<val>::minCost[n] = unable_to_insert;
                if(not findClosestKnownNodeToNode(n))
                {
                    isAbort = true;
                    return;
                }
            }
        }
    }
    MinimalSpanningForest<val>::updateMinCostAfterInsertingNode(node);
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::findClosestKnownNodeToNode(int node) {
    for(auto root: partialRoots)
    {
        if(isKnownNodeToNodeCloser(root, node))
            MinimalSpanningForest<val>::updateClosestKnownNodeOfNode(node, root);
    }
    for(auto n: MinimalSpanningForest<val>::nodes)
    {
        if(MinimalSpanningForest<val>::isKnown[n] and isKnownNodeToNodeCloser(n, node))
            MinimalSpanningForest<val>::updateClosestKnownNodeOfNode(node, n);
    }
    return MinimalSpanningForest<val>::minCost[node] != unable_to_insert;
}

template<typename val>
bool MinimalSpanningForestWithCapCons<val>::isExit(int node) {
    return isAbort or MinimalSpanningForest<val>::isExit(node);
}

template<typename val>
std::map<int, int> MinimalSpanningForestWithCapCons<val>::getDepotsToVehicleNumber() {
    std::map<int, int> depotsToVehicleNumber;
    for(auto root: partialRoots)
        depotsToVehicleNumber[root] = 0;
    for(auto node: MinimalSpanningForest<val>::nodes)
    {
        if(demandsInBranch[node])
        {
            int root = preRoot[node];
            double vehicleNumber = demandsInBranch[node] * 1.0 / vehicleCap;
            depotsToVehicleNumber[root] += ceil(vehicleNumber);
        }
    }
    return depotsToVehicleNumber;
}

template<typename val>
std::map<int, std::vector<int>> MinimalSpanningForestWithCapCons<val>::getNodesInEachRoot() {
    std::map<int, std::vector<int>> nodesInEachRoot;
    for(auto root: partialRoots)
        nodesInEachRoot[root] = std::vector<int>();
    for(auto node: MinimalSpanningForest<val>::nodes)
    {
        int root = preRoot[node];
        nodesInEachRoot[root].emplace_back(node);
    }
    return nodesInEachRoot;
}

#endif //MINIMALSPANNINGFOREST_MINIMALSPANNINGFORESTWITHCAPCONS_H
